package com.heaven.websocket;

import com.alibaba.fastjson2.JSON;
import com.heaven.constant.MessageConstants;
import com.heaven.constant.SaTokenConstants;
import com.heaven.core.entity.MessageEntity;
import com.heaven.customerServer.customerController.vo.login.LoginRes;
import com.heaven.exception.BizCodeEnum;
import com.heaven.service.LoginService;
import com.heaven.service.UserService;
import com.heaven.utils.RedisUtils;
import com.heaven.utils.UUIDUtils;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.redisson.Redisson;
import org.redisson.api.RLock;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.web.reactive.socket.WebSocketMessage;

import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * 单机处理websocket的消息转发
 * 在创建连接之前websocket的请求头必须包含 satoken:xxxxxx
 *  或者url参数携带 token=xxxxxx
 */

@Component
@Slf4j
public class StandAloneWebSocketHandle implements WebSocketHandler {
    @Resource
    private LoginService loginService;
    @Resource
    private RedisUtils redisUtils;
    @Resource
    private UserService userService;
    @Resource
    private Redisson redisson;
    @Resource
    private ThreadPoolTaskExecutor threadPoolTaskExecutor;

    /** 存储session key: userId */
    private static final Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();

    @NotNull
    @Override
    public List<String> getSubProtocols() {
        return WebSocketHandler.super.getSubProtocols();
    }

    @NotNull
    @Override
    public Mono<Void> handle(@NotNull WebSocketSession session) {
        // 握手阶段
        String token = session.getHandshakeInfo().getHeaders().getFirst(SaTokenConstants.TOKEN_HEADER_KEY);
        URI uri = session.getHandshakeInfo().getUri();
        String sessionId = session.getId();

        if(token == null) {
            // 请求头中没有携带token, 尝试从url中获取
            String[] params = uri.getQuery().split("&");
            for(String param : params) {
                String[] kv = param.split("=");
                if(kv[0].equals("token")) {
                    token = kv[1];
                    break;
                }
            }
            if(token == null) {
                log.warn("请求未携带token, url:{}, sessionId:{}", uri,sessionId);
                return session.close();
            }
        }
        LoginRes.UserSession userSession = loginService.getUserSessionByTokenWithRedis(token);
        if(userSession == null) {
            log.warn("请求携带的token无效, url:{}, sessionId:{}", uri,sessionId);
            MessageEntity resMsg = new MessageEntity("未登录");
            resMsg.setCode(BizCodeEnum.USER_NOT_LOGIN.getCode());
            writeTextToSession(session, resMsg);
            return session.close();
        }

        log.info("<<<<<< 一个连接请求建立,sessionId: {}; token: {}; uri: {} >>>>>>>>>>", sessionId, token,uri);

        return session.receive()
                .doOnSubscribe(sub -> {
                    // 连接阶段
                    sessionMap.put(userSession.getUserId(), session);
                    log.info("*********** 连接建立, sessionId: {}; 当前在线人数: {}**********", sessionId, sessionMap.size());

                })
                .flatMap(message -> {
                    // 连接后发送的数据
                    //log.info("接收到消息: {}", message);
                    switch (message.getType()) {
                        case TEXT:
                            // 处理文本消息
                            textMsgHandler(session, message);
                            break;
                        case BINARY:
                            // 二进制消息

                            break;
                        default:
                    }
                    return Mono.empty();
                })
                .then()
                .doFinally(signalType -> {
                    // 连接断开, 移除session
                    sessionMap.remove(userSession.getUserId());
                    log.info("*********** 一个连接断开, sessionId: {}; 当前在线人数: {}**********", sessionId, sessionMap.size());
                });
    }


    /**
     * 文本消息处理
     * @param session 消息源客户端session
     * @param message 消息类容
     */
    private void textMsgHandler(WebSocketSession session, WebSocketMessage message) {
        String textMsg = message.getPayloadAsText(StandardCharsets.UTF_8);
        MessageEntity msg;
        try {
            msg = JSON.to(MessageEntity.class, textMsg);
        } catch (Exception e) {
            log.error("非法消息格式: {}", e.getMessage());
            writeTextToSession(session, new MessageEntity("非法消息格式"));
            return;
        }
        if(!msg.checkIsVaild()){
            log.error("非法消息格式: {}", msg);
            writeTextToSession(session, new MessageEntity("非法消息格式"));
            return;
        }

        if(msg.getType().intValue() == MessageConstants.MESSAGE_TYPE.get("ping")
        || msg.getType().intValue() == MessageConstants.MESSAGE_TYPE.get("pong")) {
            // ping、pong消息
            writeTextToSession(session, new MessageEntity("pong", MessageConstants.MESSAGE_TYPE.get("pong")));
            return;
        }


        if(msg.getType().intValue() == MessageConstants.MESSAGE_TYPE.get("private") ||
                msg.getType().intValue() == MessageConstants.MESSAGE_TYPE.get("comment")) {
            // 私聊消息和评论消息
            WebSocketSession targetSession = getWebSocketSession(msg.getToId());
            if(targetSession == null) {
                writeTextToSession(session, new MessageEntity("目标用户不在线"));
            } else {
                writeTextToSession(targetSession, msg);
            }

            // 缓存私聊消息
            if(msg.getType().intValue() == MessageConstants.MESSAGE_TYPE.get("private")) {
                msg.setCode(0);
                msg.setStatus(0);
                /*
                 *  消息在redis中的key格式：MESSAGE_CACHE:{UUID}:{fromId}-{toId}
                 *  每个消息体都是单独缓存的，所以需要一个唯一的key
                 */
                String cacheKey = MessageConstants.MESSAGE_CACHE+ UUIDUtils.randomUUID() +":"+msg.getFromId()+"-"+msg.getToId();
                redisUtils.setCacheObject(cacheKey, msg, 1, TimeUnit.DAYS);

                CompletableFuture.runAsync(() -> {
                    // 更新用户的聊天对象ID列表
                    RLock wLock = redisson.getReadWriteLock(MessageConstants.MESSAGE_CHAT_TARGET_KEY + msg.getFromId()).writeLock();
                    wLock.lock();
                    Set<String> chatTargetIds1 = userService.selectById(msg.getFromId()).getChatTargetIds();
                    if(chatTargetIds1 == null) {
                        chatTargetIds1 = new HashSet<>();
                    }
                    if(!chatTargetIds1.contains(msg.getToId())){
                        chatTargetIds1.add(msg.getToId());
                        userService.updateOne(Query.query(Criteria.where("_id").is(msg.getFromId())), Update.update("chat_target_ids", chatTargetIds1));
                    }
                    wLock.unlock();

                    wLock = redisson.getReadWriteLock(MessageConstants.MESSAGE_CHAT_TARGET_KEY + msg.getToId()).writeLock();
                    wLock.lock();
                    Set<String> chatTargetIds2 = userService.selectById(msg.getToId()).getChatTargetIds();
                    if(chatTargetIds2 == null) {
                        chatTargetIds2 = new HashSet<>();
                    }
                    if(!chatTargetIds2.contains(msg.getFromId())){
                        chatTargetIds2.add(msg.getFromId());
                        userService.updateOne(Query.query(Criteria.where("_id").is(msg.getToId())), Update.update("chat_target_ids", chatTargetIds2));
                    }
                    wLock.unlock();
                }, threadPoolTaskExecutor);

            }
        }
    }

    /**
     * 发送文本消息给对应的客户端session
     * @param session 客户端session
     * @param message 消息内容
     */
    private void writeTextToSession(WebSocketSession session, MessageEntity message) {
        session.send(Flux.just(session.textMessage(JSON.toJSONString(message)))).subscribe();
    }

    private WebSocketSession getWebSocketSession(String userId) {
        return sessionMap.get(userId);
    }
}
